# -*- coding: utf-8 -*-
import os.path
import numpy as np

#%%
fpath = "../data/weight_mem"
BIT_W = 5
BIT_B = 16
BIT_M = 16

#%%
def write_line(file, addr, data, bit):
    assert len(data.shape) == 1
    text = 0
    for i in range(data.shape[0]):
        text |= int(data[i]) << (bit * i)
    file.write("@{:0>4X}\t{:x}\n".format(addr, text))
    return

def write_file(fname, data, bit):
    assert len(data.shape) == 2
    data[data < 0] += np.power(2, bit)
    with open(os.path.join(fpath, fname), "w") as f:
        for i in range(data.shape[0]):
            write_line(f, i, data[i], bit)
    return

def make_dwcv_mem(fname, weight, n_io):
    n_ch = weight.shape[2]
    assert len(weight.shape) == 3
    assert weight.shape[0] == weight.shape[1] == 3
    assert n_ch >= n_io and n_ch % n_io == 0
    data = weight.reshape(9, n_ch // n_io, n_io).transpose(1, 0, 2)
    assert data.shape == (n_ch // n_io, 9, n_io)
    width = n_io
    depth = n_ch // n_io * 9
    data = data.reshape(depth, width)
    
    write_file(fname, data, BIT_W)
    return

def make_pwcv_mem(fname, weight, n_in, n_out):
    n_och, n_ich = weight.shape
    assert len(weight.shape) == 2
    assert n_out % 2 == 0
    assert n_och >= n_out and n_och % n_out == 0
    assert n_ich >= n_in  and n_ich % n_in  == 0
    data = weight.reshape(n_och // n_out , n_out // 2, 2, n_ich // n_in, n_in).transpose(0, 3, 1, 4, 2)
    width = n_out * n_in
    depth = n_och // n_out * n_ich // n_in
    data = data.reshape(depth, width)
    
    write_file(fname, data, BIT_W)
    return

def make_pwcv_old(fname, weight, n_in, n_out):
    n_och, n_ich = weight.shape
    assert len(weight.shape) == 2
    assert n_och >= n_out and n_och % n_out == 0
    assert n_ich >= n_in  and n_ich % n_in  == 0
    data = weight.reshape(n_och // n_out , n_out, n_ich // n_in, n_in).transpose(0, 2, 1, 3)
    width = n_out * n_in
    depth = n_och // n_out * n_ich // n_in
    data = data.reshape(depth, width)
    
    write_file(fname, data, BIT_W)
    return

def make_norm_mem(fname, bias, mult, n_io):
    n_ch = bias.shape[0]
    assert BIT_B == BIT_M
    assert len(bias.shape) == len(mult.shape) == 1
    assert bias.shape[0] == mult.shape[0]
    assert n_ch >= n_io and n_ch % n_io == 0
    
    # data.shape == (n_ch, 2)
    data = np.hstack((
        bias.reshape(-1, 1), 
        mult.reshape(-1, 1),
    ))
    
    width = n_io * 2;
    depth = n_ch // n_io
    data = data.reshape(depth, width)
    
    write_file(fname, data, BIT_B)
    return

#%%
weight = np.load("../data/skynet_weight_int32_KKC_OCIC.npz")

make_dwcv_mem("L0_DW.mem", weight["L0-W"], 3)
make_norm_mem("L0_BM.mem", weight["L0-B"], weight["L0-M"], 1)
make_pwcv_mem("L1_PW.mem", weight["L1-W"], 3, 4)
make_norm_mem("L1_BM.mem", weight["L1-B"], weight["L1-M"], 4)

make_dwcv_mem("L2_DW.mem", weight["L2-W"], 12)
make_norm_mem("L2_BM.mem", weight["L2-B"], weight["L2-M"], 1)
make_pwcv_mem("L3_PW.mem", weight["L3-W"], 4, 6)
make_norm_mem("L3_BM.mem", weight["L3-B"], weight["L3-M"], 2)

make_dwcv_mem("L4_DW.mem", weight["L4-W"], 6)
make_norm_mem("L4_BM.mem", weight["L4-B"], weight["L4-M"], 1)
make_pwcv_mem("L5_PW.mem", weight["L5-W"], 4, 6)
make_norm_mem("L5_BM.mem", weight["L5-B"], weight["L5-M"], 1)

make_dwcv_mem("L6_DW.mem", weight["L6-W"], 3)
make_norm_mem("L6_BM.mem", weight["L6-B"], weight["L6-M"], 1)
make_pwcv_mem("L7_PW.mem", weight["L7-W"], 4, 6)
make_norm_mem("L7_BM.mem", weight["L7-B"], weight["L7-M"], 1)

make_dwcv_mem("L8_DW.mem", weight["L8-W"], 6)
make_norm_mem("L8_BM.mem", weight["L8-B"], weight["L8-M"], 1)
make_pwcv_mem("L9_PW.mem", weight["L9-W"], 8, 8)
make_norm_mem("L9_BM.mem", weight["L9-B"], weight["L9-M"], 1)

make_dwcv_mem("L10_DW.mem", weight["L10-W"], 16)
make_norm_mem("L10_BM.mem", weight["L10-B"], weight["L10-M"], 2)
make_pwcv_mem("L11_PW.mem", weight["L11-W"], 10, 4)
make_norm_mem("L11_BM.mem", weight["L11-B"], weight["L11-M"], 1)

make_pwcv_old("L1_PW_old.mem", weight["L1-W"], 3, 4)
make_pwcv_old("L12_PW_old.mem", weight["L12-W"], 2, 1)
